Pytorch --- cuda 相关

说明: 本篇是对 Pytorch Docs 0.4.1 CUDA semeantics 的部分翻译和自己在阅读过程中想加入的一些注释。

 

1、torch.cuda

　　torch.cuda 用来建立和运行 CUDA 操作。对于 Tensor 来说需要转化为 cuda tensor 来使用 cuda 计算。

2、cuda 转换

print("Outside device is 0")  # On device 0 (default in most scenarios)
with torch.cuda.device(1):
    print("Inside device is 1")  # On device 1
print("Outside device is still 0")  # On device 0

3、tensor 设置为 cuda tensor

　　三种方法：

　　a、torch.tensor（[list], device）

　　b、to(device) 实际是一种 copy() 操作而不是移动操作

　　c、cuda(device)

cuda = torch.device('cuda')     # Default CUDA device
cuda0 = torch.device('cuda:0')
cuda2 = torch.device('cuda:2')  # GPU 2 (these are 0-indexed)

x = torch.tensor([1., 2.], device=cuda0)
# x.device is device(type='cuda', index=0)
y = torch.tensor([1., 2.]).cuda()
# y.device is device(type='cuda', index=0)

with torch.cuda.device(1):
    # allocates a tensor on GPU 1
    a = torch.tensor([1., 2.], device=cuda)

    # transfers a tensor from CPU to GPU 1
    b = torch.tensor([1., 2.]).cuda()
    # a.device and b.device are device(type='cuda', index=1)

    # You can also use ``Tensor.to`` to transfer a tensor:
    b2 = torch.tensor([1., 2.]).to(device=cuda)
    # b.device and b2.device are device(type='cuda', index=1)

    c = a + b
    # c.device is device(type='cuda', index=1)

    z = x + y
    # z.device is device(type='cuda', index=0)

    # even within a context, you can specify the device
    # (or give a GPU index to the .cuda call)
    d = torch.randn(2, device=cuda2)
    e = torch.randn(2).to(cuda2)
    f = torch.randn(2).cuda(cuda2)
    # d.device, e.device, and f.device are all device(type='cuda', index=2)

　　对于上面创建的 cuda tensor 来说，已经 “绑定” 了 cuda，那么在进行加减乘除的操作时都会在其所绑定的显卡上进行。

3、torch.cuda.comm.broadcast

　　torch.cuda

　　原型：torch.cuda.comm.broadcast(tensor, devices)

　　第二个参数是元组，元组第一个 device_id 是 源 tensor 的 id。

　　示例：

cuda0 = torch.device('cuda:0')
x = troch.rand((3,4)).cuda(cuda0) # set to teh default cuda device
# the second parameter should be like (src, dst1, dst2, ...)
  xt = torch.cuda.comm.broadcast(x,(0,1)) 
# xt[0]:torch.tensor, cuda:0
# xt[1]:torch.tensor, cuda: